package com.zookeeper.lock;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

/**
 * Created by lee on 2016/11/28.
 * ZooKee分布式锁实现
 */
public class ZooKeeperLock implements Lock {

    private Log log = LogFactory.getLog(getClass());

    public final String GLOBAL_LOCK_GROUP = "globalLock";

    private final ZooKeeper zooKeeper;

    private final long minWaitMillis;
    private final long maxWaitMillis;

    private final String group;
    private final String fullGroupPath;

    private final int timeout;
    private long tryLockTime;
    private String lockName;


    private byte[] data = {0x12, 0x34};
    private int lockNum;

    public ZooKeeperLock(ZooKeeperLockConfig zooKeeperLockConfig, int timeout) {
        this(zooKeeperLockConfig, null, timeout);
    }

    public ZooKeeperLock(ZooKeeperLockConfig zooKeeperLockConfig, String group, int timeout) {
        this.group = group == null ? GLOBAL_LOCK_GROUP : group;

        this.timeout = timeout;
        this.minWaitMillis = zooKeeperLockConfig.getMinWaitMillis();
        this.maxWaitMillis = zooKeeperLockConfig.getMaxWaitMillis();

        ZooKeeperLockWatcher lockWatcher = new ZooKeeperLockWatcher();
        this.zooKeeper = ZooKeeperUtils.getZooKeeper(zooKeeperLockConfig.getConnectString(), zooKeeperLockConfig.getSessionTimeout(), lockWatcher);

        try {
            lockWatcher.awaitSyncConnected(this.timeout);

            this.fullGroupPath = zooKeeperLockConfig.getRootPath() + "/" + this.group;
            createGroupPathIfNotExists(fullGroupPath);
        } catch (Exception e) {
            throw new RuntimeException("Connection zookeeper exception!", e);
        }
    }

    private void createGroupPathIfNotExists(String groupPath) throws KeeperException, InterruptedException {
        Stat exists = zooKeeper.exists(groupPath, false);
        if(exists == null) {
            zooKeeper.create(groupPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }

    /**
     * 获取锁， 如果被占用则等待
     */
    @Override
    public void lock() {
        tryLockTime = System.currentTimeMillis();
        String prefix = "x-" + zooKeeper.getSessionId() + "-";
        String path = fullGroupPath + "/" + prefix;
        try {
            String fullLockName = zooKeeper.create(path, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            if(fullLockName == null) {
                throw new RuntimeException("Lock fair");
            }
            lockName = fullLockName.substring(fullGroupPath.length() + 1);
            lockNum = getLockNum(lockName);

            if(! checkLock(true)) {
                throw new RuntimeException("Lock fair");
            }
            if(log.isDebugEnabled()) {
                log.debug(Thread.currentThread().getName() + " Get lock [" + lockName + "] success!");
            }
        } catch (Exception e) {
            unlock();
            throw new RuntimeException("Lock exception", e);
        }
    }

    /**
     * 尝试获取锁， 不阻塞线程
     * @return true: 获取锁成功
     */
    @Override
    public boolean tryLock() {
        try {
            return checkLock(false);
        } catch (TimeoutException | KeeperException | InterruptedException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 检查锁状态
     * @param isWait 是否等待
     * @return true: 成功获取锁
     * @throws TimeoutException 锁超时
     * @throws KeeperException Zk异常
     * @throws InterruptedException 中断异常
     */
    private boolean checkLock(boolean isWait) throws TimeoutException, KeeperException, InterruptedException {
        List<String> childrens = zooKeeper.getChildren(fullGroupPath , true);
        String prevLockName = findPrevLockNum(childrens, lockNum);
        if(prevLockName == null) {
            return true;
        }

        boolean isOwnerLock;
        do {
            // 拿到锁
            isOwnerLock = isLockRelease(prevLockName);
            if(isOwnerLock || !isWait) {
                break;
            }

            // 锁超时
            if(isTimeout()) {
                throw new TimeoutException("Get lock timeout!");
            }

            // 等待获取锁
            lockWait();
        } while (true);
        return isOwnerLock;
    }

    private void lockWait() {
        long taking = System.currentTimeMillis() - tryLockTime;
        long waitMillis = (timeout - taking) / 2;
        waitMillis = Math.max(waitMillis, minWaitMillis);
        waitMillis = Math.min(waitMillis, maxWaitMillis);
        try {
            Thread.sleep(waitMillis);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    /**
     * 锁是否已释放
     * @param lockName 锁名
     * @return true: 已释放, false: 被占用
     * @throws KeeperException Zk异常
     * @throws InterruptedException 中断异常
     */
    private boolean isLockRelease(String lockName) throws KeeperException, InterruptedException {
        String lockPath = fullGroupPath + "/" + lockName;
        Stat lockStat = zooKeeper.exists(lockPath, false);
        return lockStat == null;
    }

    /**
     * 检测是否锁超时
     * @return true: 已超时
     */
    private boolean isTimeout() {
        long taking = System.currentTimeMillis() - tryLockTime;
        return taking >= timeout;
    }

    /**
     * 查询前一个lockName
     * @param childrens 包含比当前锁编号小的所有lockName
     * @param lockNum 当前锁编号
     * @return 前一个lockName
     */
    private String findPrevLockNum(List<String> childrens, int lockNum) {
        String prevLockName = null;
        int prevLockNum = -1;
        for(String children: childrens) {
            int childLockNum = getLockNum(children);
            if(childLockNum > prevLockNum && childLockNum < lockNum) {
                prevLockNum = childLockNum;
                prevLockName = children;
            }
        }
        return prevLockName;
    }

    /**
     * 释放锁
     */
    @Override
    public void unlock() {
        try {
            String path = fullGroupPath + "/" + lockName;
            Stat stat = zooKeeper.exists(path, false);
            if(stat != null) {
                if(log.isDebugEnabled()) {
                    log.debug(Thread.currentThread().getName() + " Release lock " + lockName + "");
                }
                zooKeeper.delete(path, stat.getVersion());
            }
        } catch (Exception e) {
            log.error("UnLock exception", e);
        } finally {
            try {
                zooKeeper.close();
            } catch (InterruptedException ignored) {
            }
        }
    }

    /**
     * 根据路径名获取锁编号
     * @param lockName zk路径名
     * @return 锁编号
     */
    private int getLockNum(String lockName) {
        String num = lockName.substring(lockName.lastIndexOf("-") + 1);
        return Integer.parseInt(num);
    }

    public String getGroup() {
        return group;
    }

    @Override
    public Condition newCondition() {
        throw new RuntimeException();
    }

    @Override
    public void lockInterruptibly() throws InterruptedException {
        throw new RuntimeException();
    }

    @Override
    public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
        throw new RuntimeException();
    }
}
